# !/usr/bin/env python3
# coding=utf8
"""
中国内地的(目前仅配置了期货的)(交易的)时间段配置
"""
import dataclasses
import datetime
import itertools
import pandas
from typing import Any, Dict, List, Set, Tuple, Type, Optional, Union, Callable
from vnpy.tools.chinese_mainland_business_day import ChineseMainlandBusinessDay
from vnpy.trader.constant import Exchange, Interval


class ChineseMainlandTimeConfig(object):
    """"""

    __CONFIG: Dict[tuple, List[str]] = {
        (
            ("21:00", "02:30"), ("09:00", "10:15"), ("10:30", "11:30"), ("13:30", "15:00"),
        ): [
            "SHFE.AG", "SHFE.AU",
        ],
        (
            ("21:00", "23:00"), ("09:00", "10:15"), ("10:30", "11:30"), ("13:30", "15:00"),
        ): [
            "CZCE.MA", "SHFE.RB",
        ],
        (
            ("09:30", "11:30"), ("13:00", "15:00"),
        ): [
            "CFFEX.IH", "CFFEX.IF", "CFFEX.IC", "CFFEX.IM",
        ],
    }

    CONFIG: Dict[str, List["TimeRange"]] = {}

    # 有效的日盘的开始时间和结束时间
    DAY_BEG: datetime.time = datetime.time(hour=+7, minute=0)  # +9:00 - 2
    DAY_END: datetime.time = datetime.time(hour=17, minute=0)  # 15:00 + 2
    # 有效的夜盘的开始时间和结束时间
    NIGHT_BEG: datetime.time = datetime.time(hour=19, minute=+0)  # 21:00 - 2
    NIGHT_END: datetime.time = datetime.time(hour=+4, minute=30)  # +2:30 + 2

    # https://docs.python.org/3/library/dataclasses.html
    # @dataclasses.dataclass(*, init=True, repr=True, eq=True, order=False, unsafe_hash=False, frozen=False, match_args=True, kw_only=False, slots=False, weakref_slot=False)
    @dataclasses.dataclass(frozen=True)
    class TimeRange:
        """"""
        beg: datetime.time = None  # 开始时间
        end: datetime.time = None  # 结束时间
        minutes: int = 0  # 这段时间总共有多少分钟

        def calculate(
            self,
            negative_minute: int,
            business_day: datetime.date,
            window: int,
            is_last: bool,
            next_is_day: bool,
            add_boundary_dttm: bool,
            is_default: bool,
        ) -> Tuple[List[datetime.datetime], int]:
            """
            negative_minute: 上一段时间欠了这么些分钟,
            business_day: 这段时间所处的业务日,
            window: 要按这些分钟切割出来时间段,
            is_last: 这段时间是最后一段时间,
            next_is_day: 如果为 True, 则, 还有下一段时间, 并且下一段时间是日盘,
            add_boundary_dttm: 如果这是最后一段时间, 就保存结束时间戳,
            is_default: 默认情况下(不考虑放假带来的影响, 即使放假也不丢弃夜盘),
            (这个函数没有修改类的数据成员, 可以假装这个函数是只读函数)
            前面的操作欠了 negative_minute 分钟, 所以这里要先减掉它,
            然后对剩余的时间按 window 分段, 段首的时间戳作为这一段时间的名字,
            最后可能会剩一个小的时间段, 不足 window 分钟,
            依然将其分段, 然后会欠几分钟, 让后面的时间补上, 如果后面没有时间了, 就用结束时间点终结它,
            如果没剩时间, 也要加上结束时间点,
            """
            natural_day: datetime.date = self.calculate_natural_day(business_day)

            lines: List[datetime.datetime] = []

            minutes: int = self.minutes  # 还剩这些分钟

            curr: datetime.datetime = datetime.datetime.combine(natural_day, self.beg)

            if 0 < negative_minute:
                # 有可能欠的时间超过了当前的这段时间, 所以要取两者的最小值,
                min_minute: int = min(negative_minute, minutes)
                curr += datetime.timedelta(minutes=min_minute)
                minutes -= min_minute
                negative_minute -= min_minute

            while window <= minutes:
                lines.append(curr)  # 保存这个时间段的开始时间戳
                minutes -= window
                curr += datetime.timedelta(minutes=window)

            if 0 < minutes:
                # 还剩一个小的时间段, 不足 window 分钟,
                assert minutes < window
                lines.append(curr)  # 保存这个时间段的开始时间戳
                negative_minute: int = window - minutes  # 预支了这些分钟, 让后面的时间补上,

            if is_last:  # 如果是最后一段时间了,
                curr += datetime.timedelta(minutes=minutes)
                assert curr.hour == self.end.hour and curr.minute == self.end.minute
                if add_boundary_dttm:
                    # 如果这是最后一段时间, 就保存结束时间戳,
                    lines.append(curr)  # 保存结束时间戳,
                negative_minute: int = 0

            if is_default:
                return lines, negative_minute

            if (
                True
                and not ChineseMainlandBusinessDay.has_night(business_day)
                and self.is_night()
            ):
                # 如果这个交易日没有夜盘, 并且, 这段时间是夜盘,
                if next_is_day:
                    # 如果下一段时间是日盘,
                    if 0 < negative_minute:
                        # 如果切割出来的时间段涉及了日盘,
                        if lines:
                            # 仅保留涉及日盘的那个时间段的时间戳,
                            lines = [lines[-1]]
                    else:
                        # 如果切割出来的时间段未涉及日盘, 那么, 根据这段时间切割出来的时间段是没有意义的, 应当丢弃,
                        lines.clear()
                else:
                    # 如果下一段时间也是夜盘or没有下一个时间段了, 那么, 根据这段时间切割出来的时间段是没有意义的, 应当丢弃,
                    lines.clear()

            return lines, negative_minute

        def calculate_natural_day(self, business_day: datetime.date) -> datetime.date:
            """
            提供一个交易日, 假定这个交易日有夜盘, 根据当前时间段, 计算当前时间段的开始时刻的自然日
            """
            if True:
                assert ChineseMainlandBusinessDay.is_business_day(business_day)

            if datetime.time.min <= self.beg <= ChineseMainlandTimeConfig.NIGHT_END:
                # 如果business_day是周一, 那么夜盘是上周五夜里到上周六凌晨, 这里是上周六凌晨的那一段时间,
                assert self.beg < self.end <= ChineseMainlandTimeConfig.NIGHT_END
                delta: int = 2 if business_day.weekday() == 0 else 0
                natural_day: datetime.date = business_day - datetime.timedelta(days=delta)

            elif ChineseMainlandTimeConfig.DAY_BEG <= self.beg <= ChineseMainlandTimeConfig.DAY_END:
                assert self.beg < self.end <= ChineseMainlandTimeConfig.DAY_END
                natural_day: datetime.date = business_day

            elif ChineseMainlandTimeConfig.NIGHT_BEG <= self.beg:
                # 如果business_day是周一, 那么夜盘是上周五夜里到上周六凌晨, 这里的开始时间是上周五夜里,
                delta: int = 3 if business_day.weekday() == 0 else 1
                natural_day: datetime.date = business_day - datetime.timedelta(days=delta)
                assert (
                    self.beg < self.end  # 21:00~23:00
                    or
                    self.end <= ChineseMainlandTimeConfig.NIGHT_END  # 21:00~02:30
                )
            else:
                raise RuntimeError(f"时间点异常: beg={self.beg}, end={self.end},")

            return natural_day

        def is_night(self) -> bool:
            """
            这个时间段是夜盘的时间段
            """
            if ChineseMainlandTimeConfig.DAY_BEG <= self.beg <= ChineseMainlandTimeConfig.DAY_END:
                assert self.beg < self.end <= ChineseMainlandTimeConfig.DAY_END
                return False
            return True

    @classmethod
    def get_product(cls, symbol: str) -> str:
        """
        通过'合约代码'提取出'上市品种'
        """
        if symbol.upper().endswith("L8"):
            # 特殊处理通达信行情的主力合约连续
            return symbol[:len("L8") * -1]

        for count, word in enumerate(symbol):
            if word.isdigit():
                break
        product: str = symbol[:count]
        return product

    @classmethod
    def initialize_config(cls):
        """"""
        for conf, es_product_s in cls.__CONFIG.items():
            timerange_s: List[ChineseMainlandTimeConfig.TimeRange] = []

            for item in conf:
                item: Tuple[str, str] = item
                beg_s, end_s = item
                beg_h, beg_m = beg_s.split(sep=":")
                end_h, end_m = end_s.split(sep=":")
                beg_t = datetime.time(hour=int(beg_h), minute=int(beg_m))
                end_t = datetime.time(hour=int(end_h), minute=int(end_m))

                assert beg_t != end_t

                if beg_t < end_t:  # 没跨零点(例如 21:00~23:00)
                    minutes: int = (end_t.hour * 60 + end_t.minute) - (beg_t.hour * 60 + beg_t.minute)
                else:  # 跨零点了(例如 21:00~02:30)
                    minutes: int = 24 * 60 - (beg_t.hour * 60 + beg_t.minute) + (end_t.hour * 60 + end_t.minute)

                if True:
                    _base = datetime.datetime(2020, 1, 1)
                    if beg_t < end_t:
                        _base_beg = datetime.datetime.combine(_base, beg_t)
                        _base_end = datetime.datetime.combine(_base, end_t)
                    else:
                        _base_beg = datetime.datetime.combine(_base, beg_t)
                        _base_end = datetime.datetime.combine(_base, end_t) + datetime.timedelta(days=1)
                    _delta_minutes: int = (_base_end - _base_beg).total_seconds() / 60
                    assert minutes == _delta_minutes

                timerange = ChineseMainlandTimeConfig.TimeRange(beg=beg_t, end=end_t, minutes=minutes)

                timerange_s.append(timerange)

            for es_product in es_product_s:
                cls.CONFIG[es_product] = timerange_s

    @classmethod
    def guess_business_day(cls, dttm: datetime.datetime) -> Optional[datetime.date]:
        """
        收到一个具体的日期时间点, 猜测这个时间点所处的逻辑业务日,
        只要这个时间点距离真实行情时段2小时之内, 仍认为它属于这个业务日,
        """
        business_day: datetime.date = None

        time_part: datetime.time = dttm.time()

        delta: int = 0

        # [DAY_BEG ~ DAY_END], [NIGHT_BEG ~ 24:00/00:00 ~ NIGHT_END],

        if cls.DAY_BEG <= time_part <= cls.DAY_END:
            delta: int = 0

        elif cls.NIGHT_BEG <= time_part:
            # 如果是夜盘的深夜时段,
            weekday: int = dttm.weekday()
            # 周一~周四挪1天, 周五挪3天, 周六和周日(是无效的,为了代码风格好看)挪0天,
            delta: int = 1 if weekday <= 3 else 3 if weekday == 4 else 0

        elif time_part <= cls.NIGHT_END:
            # 如果是夜盘的凌晨时段,
            weekday: int = dttm.weekday()
            # 周一(是无效的,为了代码风格好看,挪到周日)挪-1天, 周六挪2天, 周二~周五挪0天, 周日(是无效的,为了代码风格好看)挪0天,
            delta: int = -1 if weekday == 0 else 2 if weekday == 5 else 0

        else:
            # 理论上是不会走到这里的, 如果真的走的这里了, 暂时不抛异常, 而是统一取上周日的日期, 这样就不在交易日了,
            delta: int = -(dttm.weekday() + 1)  # 0是周一5是周六6是周日,

        business_day = datetime.date(year=dttm.year, month=dttm.month, day=dttm.day) + datetime.timedelta(days=delta)  # noqa E501

        if ChineseMainlandBusinessDay._in_holidays(business_day) or 5 <= business_day.weekday():
            # 交易日, 在假日表里, 或, 是周六周日,
            business_day = None

        return business_day

    @classmethod
    def calc_timepoints_for_business_day(cls, exchange: Exchange, symbol: str, business_day: datetime.date) -> Tuple[datetime.datetime, datetime.datetime]:
        """"""
        if not cls.CONFIG:
            cls.initialize_config()

        product: str = cls.get_product(symbol)

        # es => exchange+symbol 格式的数据,
        es_product: str = f"{exchange.value}.{product.upper()}"

        timerange_s: List[ChineseMainlandTimeConfig.TimeRange] = cls.CONFIG[es_product]

        tr_p0 = timerange_s[+0]  # positive_0=>p0, 正零,
        tp_p0 = datetime.datetime.combine(tr_p0.calculate_natural_day(business_day), tr_p0.beg)
        tr_n1 = timerange_s[-1]  # negative_1=>n1, 负一,
        tp_n1 = datetime.datetime.combine(tr_n1.calculate_natural_day(business_day), tr_n1.beg) + datetime.timedelta(minutes=tr_n1.minutes)  # noqa E501

        return tp_p0, tp_n1

    @classmethod
    def calc_for_business_day(
        cls,
        exchange: Exchange,
        symbol: str,
        window: int,
        is_left: bool,
        add_boundary_dttm: bool,
        business_day: datetime.date,
        is_default: bool,
    ) -> List[datetime.datetime]:
        """
        exchange: 交易所,
        symbol: 代码,
        window: 分钟间隔,
        business_day: 交易日,
        is_left: 是left,
        add_boundary_dttm: 如果left就加end如果right就加beg,
        is_default: 默认情况下(不考虑放假带来的影响)
        """
        if not cls.CONFIG:
            cls.initialize_config()

        product: str = cls.get_product(symbol)

        # es => exchange+symbol 格式的数据,
        es_product: str = f"{exchange.value}.{product.upper()}"

        timerange_s: List[ChineseMainlandTimeConfig.TimeRange] = cls.CONFIG[es_product]

        lines: List[datetime.datetime] = []

        negative_minute: int = 0

        for idx, timerange in enumerate(timerange_s):
            is_last: bool = (len(timerange_s) == idx + 1)
            # 我不是最后一个时间段, 并且, 下一个时间段是日盘,
            next_is_day: bool = not is_last and timerange_s[idx + 1].is_night() == False

            rows, negative_minute = timerange.calculate(
                negative_minute=negative_minute,
                business_day=business_day,
                window=window,
                is_last=is_last,
                next_is_day=next_is_day,
                add_boundary_dttm=True,
                is_default=is_default,
            )

            lines.extend(rows)

        if is_left and not add_boundary_dttm:
            # 如果是left规则(取左侧的时间戳)又不加边界时间戳(结束时间戳)
            lines = lines[:-1]

        if not is_left and not add_boundary_dttm:
            # 如果是right规则(取右侧的时间戳)又不加边界时间戳(开始时间戳)
            lines = lines[+1:]

        return lines

    @classmethod
    def list_chunk(cls, _src: list, size: int) -> List[list]:
        """
        直接在 python 命令行执行下面的语句, 多按两次回车, 会得到 list_chunk 函数, 它和 php 的 array_chunk 神似,
        def list_chunk(_src, size): return [_src[i:i + size] for i in range(0, len(_src), size)]
        """
        _dst: List[list] = [_src[i:i + size] for i in range(0, len(_src), size)]
        return _dst

    @classmethod
    def calc_for_business_day_deprecated(
        cls,
        exchange: Exchange,
        symbol: str,
        window: int,
        is_left: bool,
        add_boundary_dttm: bool,
        business_day: datetime.datetime,
        is_default: bool,
    ) -> List[datetime.datetime]:
        """
        C++14 的一个关键字 deprecated
        """
        if not cls.CONFIG:
            cls.initialize_config()

        product: str = cls.get_product(symbol)

        # es => exchange+symbol 格式的数据,
        es_product: str = f"{exchange.value}.{product.upper()}"

        timerange_s: List[ChineseMainlandTimeConfig.TimeRange] = cls.CONFIG[es_product]

        items: List[Tuple[datetime.datetime, bool]] = []  # ("1分钟K线"的时间戳, 这个"1分钟K线"在这个业务日是有效的)

        for idx, timerange in enumerate(timerange_s):
            last_timerange: bool = (len(timerange_s) == idx + 1)

            is_valid: bool = not (
                True
                and not ChineseMainlandBusinessDay.has_night(business_day)
                and timerange.is_night()
            )

            natural_day: datetime.datetime = timerange.calculate_natural_day(business_day=business_day)
            timerange_beg: datetime.datetime = datetime.datetime.combine(natural_day, timerange.beg)
            for offset in range(timerange.minutes + (1 if last_timerange else 0)):
                items.append((timerange_beg + datetime.timedelta(minutes=offset), is_valid))

        if window <= 1:
            # 非 is_default 时: 如果"1分钟K线"是有效的, 就取"1分钟K线"的时间戳,
            lines_chk = [_[0] for _ in items if is_default or _[1]]

        else:
            # items[-1] 是这个业务日的结束时间戳, 不是有效的"1分钟K线", 所以不参与 list_chunk 操作,
            chunk_data: List[List[Tuple[datetime.datetime, bool]]] = cls.list_chunk(items[:-1], window)

            if is_default:
                # 对于每一份数据x, 取x的首个元素的时间戳,
                lines_chk = [x[0][0] for x in chunk_data]
            else:
                # 对于每一份数据x, 只要x存在有效的元素y, 则x有效, 此时, 取x的首个元素的时间戳,
                lines_chk = [x[0][0] for x in chunk_data if [y for y in x if y[1]]]

            lines_chk.append(items[-1][0])  # 追加这个业务日的结束时间戳,

        if is_left and not add_boundary_dttm:
            # 如果是left算法(取左侧的时间戳)又不加边界时间戳(结束时间戳)
            lines_chk = lines_chk[:-1]

        if not is_left and not add_boundary_dttm:
            # 如果是right算法(取右侧的时间戳)又不加边界时间戳(开始时间戳)
            lines_chk = lines_chk[+1:]

        lines_use: List[datetime.datetime] = cls.calc_for_business_day(
            exchange=exchange,
            symbol=symbol,
            window=window,
            is_left=is_left,
            add_boundary_dttm=add_boundary_dttm,
            business_day=business_day,
            is_default=is_default,
        )

        if lines_use != lines_chk:
            for idx in range(max(len(lines_use), len(lines_chk))):
                if lines_use[idx] != lines_chk[idx]:
                    message: str = f"idx={idx}, dttm_use={lines_use[idx]}, dttm_chk={lines_chk[idx]}"
                    # print(message)
                    break
            raise Exception(
                f"exchange={exchange}, symbol={symbol}, window={window}, "
                f"is_left={is_left}, add_boundary_dttm={add_boundary_dttm}, "
                f"business_day={business_day}, is_default={is_default},"
            )

        return lines_chk

    @classmethod
    def calculate_business_days(cls, start_date: datetime.date, end_date: datetime.date) -> List[datetime.date]:
        """
        计算 [start_date, end_date] 范围的交易日, (含 start_date 和 start_date)
        """
        business_days: List[datetime.date] = []

        # date_range 函数返回的结果既包含 start 也包含 end,
        days1: pandas.DatetimeIndex = pandas.date_range(start=start_date, end=end_date, freq=pandas.DateOffset(days=1))
        days2: List[datetime.datetime] = days1.to_pydatetime().tolist()

        for day in days2:
            day: datetime.date = day.date()
            if ChineseMainlandBusinessDay.is_business_day(day):
                business_days.append(day)

        return business_days

    @classmethod
    def calc_for_days(
        cls,
        exchange: Exchange,
        symbol: str,
        window: int,
        is_left: bool,
        add_boundary_dttm: bool,
        is_default: bool,
        business_days: List[datetime.datetime] = None,
        start_date: datetime.date = None,
        end_date: datetime.date = None,
    ) -> List[datetime.datetime]:
        """
        返回window分钟K线的时间戳的列表
        """
        if business_days:
            business_days.sort()
        else:
            business_days = cls.calculate_business_days(start_date=start_date, end_date=end_date)

        lines: List[datetime.datetime] = []

        for business_day in business_days:
            rows = cls.calc_for_business_day(
                exchange=exchange,
                symbol=symbol,
                window=window,
                is_left=is_left,
                add_boundary_dttm=add_boundary_dttm,
                business_day=business_day,
                is_default=is_default,
            )
            lines.extend(rows)

        return lines

    @classmethod
    def calc_timepoints_for_days(
        cls,
        exchange: Exchange,
        symbol: str,
        start_date: datetime.date,
        end_date: datetime.date,
    ) -> Tuple[datetime.datetime, datetime.datetime]:
        """"""
        business_days = cls.calculate_business_days(start_date=start_date, end_date=end_date)

        tps: List[datetime.datetime] = []

        for business_day in business_days:
            tp_beg, tp_end = cls.calc_timepoints_for_business_day(
                exchange=exchange,
                symbol=symbol,
                business_day=business_day,
            )
            tps.append(tp_beg)
            tps.append(tp_end)

        return min(tps), max(tps)


if __name__ == "__main__":
    es_symbol_s: List[Tuple[Exchange, str]] = [
        (Exchange.SHFE, "AGL8"),
        (Exchange.SHFE, "AUL8"),
        (Exchange.CZCE, "MAL8"),
        (Exchange.CFFEX, "IHL8"),
        (Exchange.CFFEX, "IFL8"),
        (Exchange.CFFEX, "ICL8"),
        (Exchange.CFFEX, "IML8"),
    ]

    window_s: List[int] = [_ for _ in range(1, 60 * 12, 1)]

    is_left_s: List[bool] = [False, True]

    add_boundary_dttm_s: List[bool] = [False, True]

    business_day_s: List[datetime.datetime] = [
        day.date() for day in pandas.date_range(
            start=datetime.datetime(year=2022, month=1, day=1),
            end=datetime.datetime(year=2022, month=12, day=31),
            freq=pandas.DateOffset(days=1),
        ).to_pydatetime() if ChineseMainlandBusinessDay.is_business_day(day.date())
    ]

    is_default_s: List[bool] = [False, True]

    for idx, item in enumerate(
        itertools.product(
            es_symbol_s,
            window_s,
            is_left_s,
            add_boundary_dttm_s,
            business_day_s,
            is_default_s,
        )
    ):
        exchange, symbol = item[0]
        window: int = item[1]
        is_left: bool = item[2]
        add_boundary_dttm: bool = item[3]
        business_day: datetime.date = item[4]
        is_default: bool = item[5]

        ChineseMainlandTimeConfig.calc_for_business_day(
            exchange=exchange,
            symbol=symbol,
            window=window,
            is_left=is_left,
            add_boundary_dttm=add_boundary_dttm,
            business_day=business_day,
            is_default=is_default,
        )

        if idx % 1000 == 0:
            print(".", end="", flush=True)

    print("DONE, will exit...")
